#!/usr/bin/env python3
"""
最小可运行 MCP 示例（单文件）
运行方式：
    1) 先起服务端：python mcp_minimal.py server
    2) 再起客户端：python mcp_minimal.py client
"""
import json
import random
import socket
import struct
import sys

HOST, PORT = "127.0.0.1", 9876
BUF_SIZE = 4096

# ---------- 服务端 ----------
def handle_conn(conn: socket.socket):
    """一条连接 = 一个 session"""
    while True:
        # 读 4 字节长度头
        raw_len = conn.recv(4)
        if not raw_len:
            break
        msg_len = struct.unpack(">I", raw_len)[0]
        data = conn.recv(msg_len)
        if not data:
            break

        req = json.loads(data.decode())
        # 最小协议：{"type":"call","name":"<tool>","args":[...]}
        if req.get("type") == "call" and req.get("name") == "get_random":
            value = random.randint(0, 100)
            resp = json.dumps({"type": "result", "content": value}).encode()
            conn.sendall(struct.pack(">I", len(resp)) + resp)
        else:
            err = json.dumps({"type": "error", "msg": "unknown tool"}).encode()
            conn.sendall(struct.pack(">I", len(err)) + err)


def server():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind((HOST, PORT))
        s.listen(1)
        print(f"[server] listening on {HOST}:{PORT}")
        while True:
            conn, addr = s.accept()
            print(f"[server] accept {addr}")
            handle_conn(conn)

# ---------- 客户端 ----------
def call(tool_name: str, args=None):
    """极简 RPC：发一个请求，收一个响应"""
    req = json.dumps({"type": "call", "name": tool_name, "args": args or []}).encode()
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, PORT))
        s.sendall(struct.pack(">I", len(req)) + req)

        raw_len = s.recv(4)
        msg_len = struct.unpack(">I", raw_len)[0]
        resp = json.loads(s.recv(msg_len).decode())
        return resp


def client():
    print("[client] call get_random")
    ret = call("get_random")
    print("[client] received ->", ret)

# ---------- 入口 ----------
if __name__ == "__main__":
    if len(sys.argv) != 2 or sys.argv[1] not in {"server", "client"}:
        print("用法: python mcp_minimal.py server|client")
        sys.exit(1)
    if sys.argv[1] == "server":
        server()
    else:
        client()